package com.zhjw.config.db;

import lombok.extern.slf4j.Slf4j;

import java.util.Stack;

/**
 * Created by zhjw on 2023/3/6
 * <p>
 * 设置当前线程的数据源上下文，动态切换数据源
 * <p>
 * {@link DynamicRoutingDataSource} 实现Spring的AbstractRoutingDataSource接口，设置当前线程使用的数据源
 *
 * @author zhjw
 * @date 2023/3/6
 */
@Slf4j
public class DataSourceContextHolderWrapper {

    /**
     * 默认数据源
     */
    public static final String DEFAULT_DS = "master";

    private static final ThreadLocal<Stack<String>> contextHolder = new ThreadLocal<>();

    /**
     * 设置数据源名
     *
     * @param dbType DB名
     */
    public static void setDB(String dbType) {
        if (contextHolder.get() == null) {
            contextHolder.set(new Stack<>());
        }
        contextHolder.get().push(dbType);
        log.debug("Set db to {}", dbType);
    }

    /**
     * 获取数据源名
     *
     * @return DB名
     */
    public static String getDB() {
        if (contextHolder.get() != null && !contextHolder.get().empty()) {
            log.debug("Get db as {}", contextHolder.get().peek());
            return contextHolder.get().peek();
        }
        return DEFAULT_DS;
    }

    /**
     * 清除数据源名
     */
    public static void clearDB() {
        if (contextHolder.get() != null && !contextHolder.get().empty()) {
            log.debug("Pop db {}", contextHolder.get().pop());
        }
        if (contextHolder.get() != null && contextHolder.get().empty()) {
            contextHolder.remove();
        }
    }

}
